import queue
import random
import threading
import functools
import time
from concurrent.futures import Future


class State:
    RUNNING = False
    BLOCKED = True


class ThreadPoolExecutor:
    """
    A simple thread pool manager, could submit task and return immediately
    you can obtain your result from future object
    """

    class _Worker(threading.Thread):
        def __init__(self, task_queue: queue.Queue):
            """
            inner worker class will work from task_queue get task
            :param task_queue: A queue object with item (task, future)
            """
            super().__init__()
            self._task_queue = task_queue
            self._event = threading.Event()
            self.state = State.BLOCKED

        def run(self) -> None:
            """
            if not call shutdown, the run method will
            :return:
            """
            while not self._event.is_set():
                self._set_state(State.RUNNING)

                try:
                    task, future = self._task_queue.get(timeout=getattr(self, "shutdown_delay"))

                    future.set_result(task())
                    self._task_queue.task_done()
                except queue.Empty:
                    pass

                self._set_state(State.BLOCKED)

        def _set_state(self, state):
            self.state = state

        def shutdown(self):
            self._event.set()

        def __bool__(self):
            return self.state

    def __init__(self, max_worker, shutdown_delay=2):
        self.max = max_worker
        self._task_queue = queue.Queue()
        self.workers = []
        self._Worker.shutdown_delay = shutdown_delay

    def submit(self, fn, *args, **kwargs):
        """
        async submit a function to thread pool and return future object
        you could get fn result from future object
        :param fn: 
        :param args: 
        :param kwargs: 
        :return: 
        """
        self._has_worker()
        future = Future()
        task = functools.partial(fn, *args, **kwargs)
        self._task_queue.put((task, future))
        return future

    def _has_worker(self):
        """
        will judge have worker exit, if have and blocked state will not create new thread
        if not workers or workers less than max workers will create new thread, first, there
        are no blocked worker
        :return:
        """
        if list(filter(None, self.workers)):
            return

        if not self.workers or len(self.workers) < self.max:
            # create new workers
            worker = self._Worker(self._task_queue)
            self.workers.append(worker)
            worker.start()

    def shutdown(self):
        """
        send signal to all workers to shutdown
        :return:
        """
        for worker in self.workers:
            worker.shutdown()


if __name__ == '__main__':
    def test_func(n=100):
        for i in range(n):
            print(i)
        return random.randrange(100)


    executor = ThreadPoolExecutor(10, shutdown_delay=1)
    fu = executor.submit(test_func, 1)
    executor.submit(test_func, 100)
    executor.submit(test_func, 10)
    executor.submit(test_func, 10)

    time.sleep(2)
    print(fu.result())
    executor.shutdown()

